//
// Created by bobin on 17-11-16.
//

#include "G2OTypes.h"
#include "Eigen/Dense"
#include "opencv2/core/core.hpp"
#include "vector"
#include <fstream>
#include "Optimize.h"
#include "string"
#include "chrono"

struct PointData {

    PointData();

    void plane_points();

    void space_points();

    void line_points();

    void project_points(double f, double cx, double cy, Eigen::Matrix3d R, Eigen::Vector3d t);

    void DumpFile(std::string fileName);

    std::vector<Eigen::Vector3d> x3dpt;
    std::vector<Eigen::Vector2d> x2dpt;

    std::vector<Eigen::Vector4d> planes;
    std::vector<g2o::Vector6D> lines;

    std::vector<int> pointCate;
    std::vector<int> lineIndex;
    std::vector<int> planeIndex;


    const int point_in_plane_num = 100;
    const int point_in_space_num = 100;
    const int point_in_line_num = 100;
};

PointData::PointData() {
    int all_point_num = point_in_line_num + point_in_space_num + point_in_plane_num;
    x3dpt.resize(all_point_num);
    x2dpt.resize(all_point_num);
    planeIndex.resize(all_point_num);
    lineIndex.resize(all_point_num);
    pointCate.resize(all_point_num);
}

void PointData::plane_points() {
    Eigen::Vector4d plane;
    plane << 1, 0, 0, 0;
    planes.push_back(plane);
    plane << 0, 0, 1, 0;
    planes.push_back(plane);
    int plane_start = point_in_space_num + point_in_line_num;
    for (int j = plane_start; j < plane_start + point_in_plane_num; ++j) {
        Eigen::Vector3d pt;
        int x = j - plane_start;
        if (x < point_in_plane_num / 2) {
            pt << 0, x / 10, x % 10;
            planeIndex[j] = 0;
        } else {
            pt << j / 10, j % 10, 0;
            planeIndex[j] = 1;
        }
        x3dpt[j] = pt;
        pointCate[j] = Optimize::PlanePoint;
        lineIndex[j] = -1;
    }

}

void PointData::space_points() {

    for (int j = 0; j < point_in_space_num; ++j) {
        Eigen::Vector3d pt;
        pt << j / 20, (j % 20) / 4, (j % 20) % 4;
        x3dpt[j] = pt;
        pointCate[j] = Optimize::SpacePoint;
        lineIndex[j] = -1;
        planeIndex[j] = -1;
    }
}

void PointData::line_points() {
    int point_start = point_in_space_num;
    for (int i = point_start; i < point_start + point_in_line_num; ++i) {
        Eigen::Vector3d pt;
        double x, y, z;
        x = i - point_start;
        if (x < point_in_line_num / 3) {
            y = 100 - x;
            z = 2 * y;
            lineIndex[i] = 0;
        } else if (x < point_in_line_num * 2 / 3) {
            y = 120 - x;
            z = 2 * y;
            lineIndex[i] = 1;
        } else {
            y = 130 - x;
            z = 2 * y;
            lineIndex[i] = 2;
        }
        pt << x, y, z;
        x3dpt[i] = pt;
        pointCate[i] = Optimize::LinePoint;
        planeIndex[i] = -1;
    }
    g2o::Vector6D line;
    line << 0, 100, 200, 99, 1, 2;
    lines.push_back(line);

    line << 0, 120, 240, 99, 21, 42;
    lines.push_back(line);
    line << 0, 130, 260, 1, 31, 62;
    lines.push_back(line);
}

void PointData::project_points(double f, double cx, double cy, Eigen::Matrix3d R, Eigen::Vector3d t) {
    float noise_pixel = 1;
    auto get_rand = [noise_pixel]() { return noise_pixel * ((rand() % 1001 - 500) / 5.0); };
    for (int i = 0; i < x3dpt.size(); ++i) {
        Eigen::Vector3d ptc;
        ptc = R * x3dpt[i] + t;
//        std::cout << ptc.transpose() << std::endl;
        Eigen::Vector2d uv;
        uv[0] = ptc[0] / ptc[2] * f + cx + get_rand();
        uv[1] = ptc[1] / ptc[2] * f + cy + get_rand();
        x2dpt[i] = uv;
    }
}

void PointData::DumpFile(std::string fileName) {
    std::ofstream ofs;
    ofs.open(fileName, std::ios_base::out);
    for (int i = 0; i < x3dpt.size(); ++i) {
        ofs << x3dpt[i].transpose() << '\t' << pointCate[i] << '\t' << lineIndex[i]
            << '\t' << planeIndex[i] << std::endl;
    }

    for (auto pt : x2dpt) {
        ofs << pt.transpose() << std::endl;
    }
    ofs.close();
}

int main(int argc, char **argv) {
    std::cout << "test plane line ba" << std::endl;
    PointData pointData;
    pointData.space_points();
    std::cout << "3d point number is : " << pointData.x3dpt.size() << std::endl;
    pointData.line_points();
    std::cout << "3d point number is : " << pointData.x3dpt.size() << pointData.lineIndex.size() << std::endl;
    pointData.plane_points();
    std::cout << "3d point number is : " << pointData.x3dpt.size() << pointData.planeIndex.size() << std::endl;


    double f = 800;
    double cx = 320;
    double cy = 240;

    Eigen::Matrix3d R = (Eigen::AngleAxisd(0.5, Eigen::Vector3d::UnitZ()) *
                         Eigen::AngleAxisd(-0.5, Eigen::Vector3d::UnitY()) *
                         Eigen::AngleAxisd(0.5, Eigen::Vector3d::UnitX())
    ).matrix();
    Eigen::Vector3d t(1, -2, 10);

    Eigen::Matrix3d Rcw = R.transpose();
    Eigen::Vector3d tcw = -Rcw * t;

    std::cout << "gt rot: " << std::endl;
    std::cout << R << std::endl;

    std::cout << "gt trans" << std::endl;
    std::cout << t << std::endl;
    pointData.project_points(f, cx, cy, R, t);

    pointData.DumpFile("point.txt");

    Eigen::Matrix3d K;
    K << f, 0, cx, 0, f, cy, 0, 0, 1;
    Eigen::Matrix3d rot;
    Eigen::Vector3d trans;
    Optimize optimizer(pointData.x3dpt, pointData.x2dpt, K, pointData.pointCate, pointData.lineIndex,
                       pointData.planeIndex, pointData.lines,
                       pointData.planes);

    std::chrono::steady_clock::time_point t1 = std::chrono::steady_clock::now();
    optimizer.bundleAdjust(rot, trans, true);
    std::chrono::steady_clock::time_point t2 = std::chrono::steady_clock::now();
    double timeCost1 = std::chrono::duration_cast<std::chrono::duration<double> >(t2 - t1).count();
    std::cout << "optimize with structure cost time: " << timeCost1 << std::endl;
    std::cout << "rot: " << std::endl;
    std::cout << rot << std::endl;
    std::cout << "trans" << std::endl;
    std::cout << trans << std::endl;


    optimizer.bundleAdjust(rot, trans);
    std::chrono::steady_clock::time_point t3 = std::chrono::steady_clock::now();
    double timeCost2 = std::chrono::duration_cast<std::chrono::duration<double> >(t3 - t2).count();
    std::cout << "optimize without structure cost time: " << timeCost2 << std::endl;
    std::cout << "rot: " << std::endl;
    std::cout << rot << std::endl;
    std::cout << "trans" << std::endl;
    std::cout << trans << std::endl;
    return 0;
}
